import matplotlib.pyplot as plt 
import pandas as pd
import seaborn as sns
import sys 
import math


# def naive_graph(data, width=None, part: int=25):
#     if not width:
#         width = (max(data) - min(data)) / 25
#     bins = range(min(data), max(data) + width, width)  # 从最小值到最大值，每5个单位一个区间

#     # 使用cut函数将数据分组
#     labels = ['{:d}-{:d}'.format(bins[i], bins[i+1]-1) for i in range(len(bins)-1)]
#     grouped = pd.cut(data, bins=bins, labels=labels, right=False)

#     # 统计每个区间的数量
#     value_counts = grouped.value_counts().sort_index()

#     # 打印统计结果
#     # print(value_counts)

#     # 画出统计图
#     value_counts.plot(kind='bar')
#     plt.title('Value Counts by Interval')
#     plt.xlabel('Value Interval')
#     plt.ylabel('Count')
#     plt.show()
    
    
# def scatter(x, y, labels=None):

#     # 准备数据
#     if not labels:
#         labels = ['A'] * len(x) # 点的标签
#     if len(x) != len(y) and len(x) != len(labels):
#         raise ValueError("Length of x and y and labels are not equal.")

#     # 创建点图
#     plt.scatter(x, y)

#     # 为每个点添加标签
#     for i, txt in enumerate(labels):
#         plt.annotate(txt, (x[i], y[i]))

#     # 添加标题和轴标签
#     plt.title('Simple Scatter Plot')
#     plt.xlabel('X Axis')
#     plt.ylabel('Y Axis')

#     # 显示图表
#     plt.show()
    
    
# def heatmap(matrix):
#     sns.heatmap(matrix, annot=True, cmap='coolwarm')
#     plt.title('Heatmap Example')
#     plt.show()


def draw_length(datas: list[str], 
           sort: bool=False,
           reverse: bool=False,
           tokenizer_path: str='/mnt/public/open_source_model/Qwen2.5/Qwen2.5-3B',
        #    use_tqdm: bool=True,
           ) -> None:

    from transformers import AutoTokenizer 
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    
    tokenized = []
    n=8
    group = math.ceil(len(datas)/n)
    datas_list = [datas[i*n: (i+1)*n] for i in range(group)]
    for sub_datas in datas_list:
        tokenized.extend(tokenizer.encode(sub_datas))
    data_length = [len(i) for i in tokenized]
    
    if sort:
        data_length = sorted(data_length, reverse=reverse)
    
    plt.figure()
    plt.plot([i for i in range(len(data_length))], data_length, marker='o')
    plt.grid(True)
    plt.show()


if __name__ == '__main__':
    draw_length(['alsdjf','aodnslv','zodljfoiwe','alsdjf','aodnslv','zodljfoiwe','alsdjf','aodnslv','zodljfoiwe','alsdjf','aodnslv','zodljfoiwe','alsdjf','aodnslv','zodljfoiwe'])